﻿#include"yolo.h"

using namespace std;
using namespace cv;
using namespace cv::dnn;

bool Yolo::readModel(Net& net, string& netPath, bool isCuda = false) {
    try {
        net = readNet(netPath);
    }
    catch (const std::exception&) {
        return false;
    }
    //cuda
    if (isCuda) {
        net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
        net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA);//_FP16
    }
    //cpu
    else {
        net.setPreferableBackend(cv::dnn::DNN_BACKEND_DEFAULT);
        net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
    }
    return true;
}

bool Yolo::Detect(Mat& SrcImg, Net& net, vector<Output>& output,std::vector<std::string>& className) {
    static Mat blob;
    int col = SrcImg.cols;
    int row = SrcImg.rows;
    int maxLen = MAX(col, row);
    //补边成正方形，否则会因为拉伸成正方形导致识别率降低
    Mat netInputImg = SrcImg.clone();
    if (maxLen > 1.2 * col || maxLen > 1.2 * row) {
        Mat resizeImg = Mat::zeros(maxLen, maxLen, CV_8UC3);
        SrcImg.copyTo(resizeImg(Rect(0, 0, col, row)));
        netInputImg = resizeImg;
    }
    blobFromImage(netInputImg, blob, 1 / 255.0, cv::Size(netWidth, netHeight), cv::Scalar(0, 0, 0), true, false);
    // If there is no problem with other settings but the result deviates greatly , Try the following two sentences //
    //blobFromImage(netInputImg, blob, 1 / 255.0, cv::Size(netWidth, netHeight), cv::Scalar(104, 117, 123), true, false);
    //blobFromImage(netInputImg, blob, 1 / 255.0, cv::Size(netWidth, netHeight), cv::Scalar(114, 114,114), true, false);
    net.setInput(blob);
    static std::vector<cv::Mat> netOutputImg;
    netOutputImg.clear();
    //vector<string> outputLayerName{"345","403", "461","output" };
    //net.forward(netOutputImg, outputLayerName[3]); // obtain output Output //
    try
    { //release OK
        net.forward(netOutputImg, net.getUnconnectedOutLayersNames());//debug Report errors initCUDABackend CUDA backend will fallback to the CPU implementation for the layer "_input"
    }
    catch (const std::exception& e)
    {
        cout << e.what();
    }
    static std::vector<int> classIds;// result id Array //
    static std::vector<float> confidences;// Results each id Corresponding confidence array //
    static std::vector<cv::Rect> boxes;// Every id Rectangle box //
    classIds.clear();
    confidences.clear();
    boxes.clear();
    float ratio_h = (float)netInputImg.rows / netHeight;
    float ratio_w = (float)netInputImg.cols / netWidth;
    int net_width = className.size() + 5; // The output network width is the number of categories +5//
    float* pdata = (float*)netOutputImg[0].data;
    for (int stride = 0; stride < strideSize; stride++) { //stride
        int grid_x = (int)(netWidth / netStride[stride]);
        int grid_y = (int)(netHeight / netStride[stride]);
        for (int anchor = 0; anchor < 3; anchor++) { //anchors
            const float anchor_w = netAnchors[stride][anchor * 2];
            const float anchor_h = netAnchors[stride][anchor * 2 + 1];
            for (int i = 0; i < grid_y; i++) {
                for (int j = 0; j < grid_x; j++) {
                    float box_score = pdata[4]; ;// Get... For each row box The probability that the box contains an object //
                    if (box_score >= boxThreshold) {
                        cv::Mat scores(1, className.size(), CV_32FC1, pdata + 5);
                        Point classIdPoint;
                        double max_class_socre;
                        minMaxLoc(scores, 0, &max_class_socre, 0, &classIdPoint);
                        max_class_socre = (float)max_class_socre;
                        if (max_class_socre >= classThreshold) {
                            //rect [x,y,w,h]
                            float x = pdata[0]; //x
                            float y = pdata[1]; //y
                            float w = pdata[2]; //w
                            float h = pdata[3]; //h
                            int left = (x - 0.5 * w) * ratio_w;
                            int top = (y - 0.5 * h) * ratio_h;
                            classIds.push_back(classIdPoint.x);
                            confidences.push_back(max_class_socre * box_score);
                            boxes.push_back(Rect(left, top, int(w * ratio_w), int(h * ratio_h)));
                        }
                    }
                    pdata += net_width;// The next line //
                }
            }
        }
    }
    // Non maximum suppression is performed to eliminate redundant overlapping boxes with low confidence （NMS）//
    static vector<int> nms_result;
    nms_result.clear();
    NMSBoxes(boxes, confidences, nmsScoreThreshold, nmsThreshold, nms_result);
    Output result;
    for (int i = 0; i < nms_result.size(); i++) {
        int idx = nms_result[i];
        result.id = classIds[idx];
        result.confidence = confidences[idx];
        result.box = boxes[idx];
        output.push_back(result);
    }
    if (output.size())
        return true;
    else
        return false;
}

void Yolo::drawPred(Mat& img, vector<Output> result, vector<Scalar> color,std::vector<std::string>& className) {
    for (int i = 0; i < result.size(); i++) {
        int left, top;
        left = result[i].box.x;
        top = result[i].box.y;
        int color_num = i;
        //rectangle(img, result[i].box, color[result[i].id], 2, 8);
        rectangle(img, result[i].box, color[result[i].id], 2, 8);
        string label = className[result[i].id] + ":" + to_string(result[i].confidence);
        int baseLine;
        Size labelSize = getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
        top = max(top, labelSize.height);
        //rectangle(frame, Point(left, top - int(1.5 * labelSize.height)), Point(left + int(1.5 * labelSize.width), top + baseLine), Scalar(0, 255, 0), FILLED);
        putText(img, label, Point(left, top), FONT_HERSHEY_SIMPLEX, 1, color[result[i].id], 2);
    }
    //imshow("1", img);
    imwrite("out.bmp", img);
    //waitKey();
    //destroyAllWindows();
}
